import jax
from jax import random
import flax
from flax import linen as nn
from typing import Sequence


class ExplicitMLP(nn.Module):
    features: Sequence[int]

    def setup(self) -> None:
        self.layers = [nn.Dense(feat) for feat in self.features]

    def __call__(self, inputs):
        x = inputs
        for i, lyr in enumerate(self.layers):
            x = lyr(x)
            if i != len(self.layers) - 1:
                x = nn.relu(x)
        return x

key = random.PRNGKey(0)
key1, key2 = random.split(key, 2)

x = random.uniform(key1, (4, 4))
model = ExplicitMLP(features=[3, 4, 5])
params = model.init(key2, x)
# y = model(x)  # AttributeError: "ExplicitMLP" object has no attribute "layers". If "layers" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'.
y = model.apply(params, x)

print('init params', jax.tree_util.tree_map(lambda x: x.shape, params))
print('x', x)
print('y', y)
